# 直接回归Tschebyscheff系数
import argparse
from typing import Dict
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from rsp.antenna.array import linear_array
from apps.wfs.wfs_config import WfsConfig as WG

class TcrMlp(nn.Module):
    def __init__(self, N:int = 8):
        '''
        N为线性阵列天线个数，输入层：天线数量、频率、间距、扫描角度（角度）
        '''
        super().__init__()
        self.l1_fc = nn.Linear(4, 32, dtype=torch.float64)
        self.l1_relu = nn.ReLU()
        self.l2_fc = nn.Linear(32, 128, dtype=torch.float64)
        self.l2_relu = nn.ReLU()
        self.l3_fc = nn.Linear(128, 512, dtype=torch.float64)
        self.l3_relu = nn.ReLU()
        self.l4_fc = nn.Linear(512, 128, dtype=torch.float64)
        self.l4_relu = nn.ReLU()
        self.l5_fc = nn.Linear(128, 32, dtype=torch.float64)
        self.l5_relu = nn.ReLU()
        self.l6_fc = nn.Linear(32, N, dtype=torch.float64)
        self.y_fn = nn.Sigmoid()

    def forward(self, x):
        a1 = self.l1_relu(self.l1_fc(x))
        a2 = self.l2_relu(self.l2_fc(a1))
        a3 = self.l3_relu(self.l3_fc(a2))
        a4 = self.l4_relu(self.l4_fc(a3))
        a5 = self.l5_relu(self.l5_fc(a4))
        return self.l6_fc(a5)